# Copyright (c) SenseTime. All Rights Reserved.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import numpy as np
import torch.nn.functional as F
from pysot.core.config import cfg
from pysot.utils.anchor import Anchors
from pysot.tracker.base_tracker import SiameseTracker
from socialgan.test_model01 import general_pred


class SiamRPNTracker(SiameseTracker):
    def __init__(self, model):
        super(SiamRPNTracker, self).__init__()
        self.score_size = (cfg.TRACK.INSTANCE_SIZE - cfg.TRACK.EXEMPLAR_SIZE) // \
            cfg.ANCHOR.STRIDE + 1 + cfg.TRACK.BASE_SIZE
        self.anchor_num = len(cfg.ANCHOR.RATIOS) * len(cfg.ANCHOR.SCALES)
        hanning = np.hanning(self.score_size)
        window = np.outer(hanning, hanning)
        self.window = np.tile(window.flatten(), self.anchor_num)
        self.anchors = self.generate_anchor(self.score_size)
        self.model = model
        self.model.eval()

    def generate_anchor(self, score_size):
        anchors = Anchors(cfg.ANCHOR.STRIDE,
                          cfg.ANCHOR.RATIOS,
                          cfg.ANCHOR.SCALES)
        anchor = anchors.anchors
        x1, y1, x2, y2 = anchor[:, 0], anchor[:, 1], anchor[:, 2], anchor[:, 3]
        anchor = np.stack([(x1+x2)*0.5, (y1+y2)*0.5, x2-x1, y2-y1], 1)
        total_stride = anchors.stride
        anchor_num = anchor.shape[0]
        anchor = np.tile(anchor, score_size * score_size).reshape((-1, 4))
        ori = - (score_size // 2) * total_stride
        xx, yy = np.meshgrid([ori + total_stride * dx for dx in range(score_size)],
                             [ori + total_stride * dy for dy in range(score_size)])
        xx, yy = np.tile(xx.flatten(), (anchor_num, 1)).flatten(), \
            np.tile(yy.flatten(), (anchor_num, 1)).flatten()
        anchor[:, 0], anchor[:, 1] = xx.astype(np.float32), yy.astype(np.float32)
        return anchor

    def _convert_bbox(self, delta, anchor):
        delta = delta.permute(1, 2, 3, 0).contiguous().view(4, -1)
        delta = delta.data.cpu().numpy()

        delta[0, :] = delta[0, :] * anchor[:, 2] + anchor[:, 0]
        delta[1, :] = delta[1, :] * anchor[:, 3] + anchor[:, 1]
        delta[2, :] = np.exp(delta[2, :]) * anchor[:, 2]
        delta[3, :] = np.exp(delta[3, :]) * anchor[:, 3]
        return delta

    def _convert_score(self, score):
        score = score.permute(1, 2, 3, 0).contiguous().view(2, -1).permute(1, 0)
        score = F.softmax(score, dim=1).data[:, 1].cpu().numpy()
        return score

    def _bbox_clip(self, cx, cy, width, height, boundary):
        cx = max(0, min(cx, boundary[1]))
        cy = max(0, min(cy, boundary[0]))
        width = max(10, min(width, boundary[1]))
        height = max(10, min(height, boundary[0]))
        return cx, cy, width, height

    def init(self, img, bbox):
        """
        args:
            img(np.ndarray): BGR image
            bbox: (x, y, w, h) bbox
        """
        self.center_pos = np.array([bbox[0]+(bbox[2]-1)/2,
                                    bbox[1]+(bbox[3]-1)/2])
        self.size = np.array([bbox[2], bbox[3]])

        # calculate z crop size
        w_z = self.size[0] + cfg.TRACK.CONTEXT_AMOUNT * np.sum(self.size)
        h_z = self.size[1] + cfg.TRACK.CONTEXT_AMOUNT * np.sum(self.size)
        s_z = round(np.sqrt(w_z * h_z))

        # calculate channle average
        self.channel_average = np.mean(img, axis=(0, 1))

        # get crop
        z_crop = self.get_subwindow(img, self.center_pos,
                                    cfg.TRACK.EXEMPLAR_SIZE,
                                    s_z, self.channel_average)
        self.model.template(z_crop)


    def track(self, img, fram_order, all_bbox):
        """
        args:
            img(np.ndarray): BGR image
        return:
            bbox(list):[x, y, width, height]
        """

        w_z = self.size[0] + cfg.TRACK.CONTEXT_AMOUNT * np.sum(self.size)
        h_z = self.size[1] + cfg.TRACK.CONTEXT_AMOUNT * np.sum(self.size)
        s_z = np.sqrt(w_z * h_z)
        scale_z = cfg.TRACK.EXEMPLAR_SIZE / s_z
        s_x = s_z * (cfg.TRACK.INSTANCE_SIZE / cfg.TRACK.EXEMPLAR_SIZE)
        x_crop = self.get_subwindow(img, self.center_pos,
                                    cfg.TRACK.INSTANCE_SIZE,
                                    round(s_x), self.channel_average)

        outputs = self.model.track(x_crop)

        score = self._convert_score(outputs['cls'])
        pred_bbox = self._convert_bbox(outputs['loc'], self.anchors)
        peak_distance = outputs['distance']

        def change(r):
            return np.maximum(r, 1. / r)

        def sz(w, h):
            pad = (w + h) * 0.5
            return np.sqrt((w + pad) * (h + pad))

        # scale penalty
        s_c = change(sz(pred_bbox[2, :], pred_bbox[3, :]) /
                     (sz(self.size[0]*scale_z, self.size[1]*scale_z)))

        # aspect ratio penalty
        r_c = change((self.size[0]/self.size[1]) /
                     (pred_bbox[2, :]/pred_bbox[3, :]))
        penalty = np.exp(-(r_c * s_c - 1) * cfg.TRACK.PENALTY_K)
        pscore = penalty * score

        # window penalty
        pscore = pscore * (1 - cfg.TRACK.WINDOW_INFLUENCE) + \
            self.window * cfg.TRACK.WINDOW_INFLUENCE
        best_idx = np.argmax(pscore)


        global obs_traj, block_fram, block, len
        # obs_traj = []

        if fram_order == 2:  # demo
        # if fram_order == 1:  # test
            block = 0

        if isinstance(peak_distance, int):
            min_peak_distance = peak_distance
        elif isinstance(peak_distance, float):
            min_peak_distance = peak_distance
        else:
            min_peak_distance = min(peak_distance)

        all_len = len(all_bbox)
        i = 0.8

        threshold = i*(score[best_idx]/0.95) + (1-i)*(min_peak_distance/5.5)
        # threshold = i*((score[best_idx]-0.95)/(1-0.95)) + (1-i)*((min_peak_distance-5.5)/(10-5.5))

        # if (fram_order == 283 and block == 0 and all_len >= 18) or (fram_order == 296 and block == 0 and all_len >= 18):
        # if (fram_order == 282 or fram_order == 295 or fram_order == 307) and block == 0 and all_len >= 18:

        if score[best_idx] < 0 and block == 0 and all_len >= 18:
        # if score[best_idx] < 0.7 and block == 0 and all_len >= 18:
        # if threshold < 0.851 and block == 0 and all_len > 18:
        # if min_peak_distance <= 3 and block == 0 and all_len >= 18:

        # global obs_traj, block_fram, block, len

        # if score[best_idx] < 0.75 and min_peak_distance <= 3.2 and block == 0 and all_len >= 18:
        # if score[best_idx] < 0.99892957 and score[best_idx] > 0.99892955 and block == 0 and all_len >= 18:

            print('*********************** block ***********************')

            block_fram = 0
            block = 1
            obs_traj = all_bbox[fram_order - 19:fram_order]


            bbox = pred_bbox[:, best_idx] / scale_z
            lr = penalty[best_idx] * score[best_idx] * cfg.TRACK.LR

            cx = bbox[0] + self.center_pos[0]
            cy = bbox[1] + self.center_pos[1]

            # smooth bbox
            width = self.size[0] * (1 - lr) + bbox[2] * lr
            height = self.size[1] * (1 - lr) + bbox[3] * lr

            # clip boundary
            cx, cy, width, height = self._bbox_clip(cx, cy, width, height, img.shape[:2])
            
            # udpate state
            self.center_pos = np.array([cx, cy])
            self.size = np.array([width, height])
            bbox = [cx - width / 2, cy - height / 2,width, height]
            best_score = score[best_idx]

            print('judge', 'bbox{}:'.format(fram_order), bbox)

        # elif score[best_idx] < 0 and block == 0 and all_len >= 18:
        # elif score[best_idx] < 0.3 and block == 0 and all_len >= 18:
        # # elif min_peak_distance <= 3 and block == 0 and all_len >= 18:
        # # elif score[best_idx] < 0.88 and min_peak_distance <= 4 and block == 0 and all_len >= 18:
        #     # if score[best_idx] < 0.99892957 and score[best_idx] > 0.99892955 and block == 0 and all_len >= 18:
        #
        #     # threshold = i*((score[best_idx]-0.95)/(1-0.95)) + (1-i)*((min_peak_distance-5.5)/(10-5.5))
        #     # if threshold < 2 and block == 0 and all_len > 18:
        #     # if min_peak_distance < 3.01 and block == 0 and all_len > 18:
        #     # if min_peak_distance < 3.01 and score[best_idx] < 0.99 and block == 0 and all_len > 18:
        #     print('*********************** block ***********************')
        #
        #     block_fram = 0
        #     block = 1
        #     obs_traj = all_bbox[fram_order - 19:fram_order]
        #
        #     bbox = pred_bbox[:, best_idx] / scale_z
        #     lr = penalty[best_idx] * score[best_idx] * cfg.TRACK.LR
        #
        #     cx = bbox[0] + self.center_pos[0]
        #     cy = bbox[1] + self.center_pos[1]
        #
        #     # smooth bbox
        #     width = self.size[0] * (1 - lr) + bbox[2] * lr
        #     height = self.size[1] * (1 - lr) + bbox[3] * lr
        #
        #     # clip boundary
        #     cx, cy, width, height = self._bbox_clip(cx, cy, width, height, img.shape[:2])
        #
        #     # udpate state
        #     self.center_pos = np.array([cx, cy])
        #     self.size = np.array([width, height])
        #     bbox = [cx - width / 2, cy - height / 2, width, height]
        #     best_score = score[best_idx]
        #
        #     print('judge', 'bbox{}:'.format(fram_order), bbox)

        # elif score[best_idx] < 0.99999 and block == 0 and all_len == 18:
        #     # if min_peak_distance < 5.5 and block == 0 and all_len > 18:
        #     print('*********************** 18**block ***********************')
        #
        #     block_fram = 0
        #     block = 1
        #     obs_traj = all_bbox[fram_order - 19:fram_order]
        #
        #     bbox = pred_bbox[:, best_idx] / scale_z
        #     lr = penalty[best_idx] * score[best_idx] * cfg.TRACK.LR
        #
        #     cx = bbox[0] + self.center_pos[0]
        #     cy = bbox[1] + self.center_pos[1]
        #
        #     # smooth bbox
        #     width = self.size[0] * (1 - lr) + bbox[2] * lr
        #     height = self.size[1] * (1 - lr) + bbox[3] * lr
        #
        #     # clip boundary
        #     cx, cy, width, height = self._bbox_clip(cx, cy, width, height, img.shape[:2])
        #
        #     # udpate state
        #     self.center_pos = np.array([cx, cy])
        #     self.size = np.array([width, height])
        #     bbox = [cx - width / 2, cy - height / 2, width, height]
        #     best_score = score[best_idx]
        #
        #     print('judge', 'bbox{}:'.format(fram_order), bbox)

        # # elif threshold < 2 and block == 0 and all_len == 10:
        # elif score[best_idx] < 0.95 and block == 0 and all_len == 10:
        # # elif min_peak_distance < 5.5 and block == 0 and all_len == 10:
        #     print('*********************** 9 ** block ***********************')
        #     block_fram = 0
        #     block = 1
        #
        #     obs_traj = [all_bbox[0],
        #                 [(all_bbox[0][0] + all_bbox[1][0])/2,(all_bbox[0][1] + all_bbox[1][1])/2,
        #                  (all_bbox[0][2] + all_bbox[1][2])/2,(all_bbox[0][3] + all_bbox[1][3])/2],
        #                 all_bbox[1],
        #                 [(all_bbox[1][0] + all_bbox[2][0]) / 2, (all_bbox[1][1] + all_bbox[2][1]) / 2,
        #                  (all_bbox[1][2] + all_bbox[2][2]) / 2, (all_bbox[1][3] + all_bbox[2][3]) / 2],
        #                 all_bbox[2],
        #                 [(all_bbox[2][0] + all_bbox[3][0]) / 2, (all_bbox[2][1] + all_bbox[3][1]) / 2,
        #                  (all_bbox[2][2] + all_bbox[3][2]) / 2, (all_bbox[2][3] + all_bbox[3][3]) / 2],
        #                 all_bbox[3],
        #                 [(all_bbox[3][0] + all_bbox[4][0]) / 2, (all_bbox[3][1] + all_bbox[4][1]) / 2,
        #                  (all_bbox[3][2] + all_bbox[4][2]) / 2, (all_bbox[3][3] + all_bbox[4][3]) / 2],
        #                 all_bbox[4],
        #                 [(all_bbox[4][0] + all_bbox[5][0]) / 2, (all_bbox[4][1] + all_bbox[5][1]) / 2,
        #                  (all_bbox[4][2] + all_bbox[5][2]) / 2, (all_bbox[4][3] + all_bbox[5][3]) / 2],
        #                 all_bbox[5],
        #                 [(all_bbox[5][0] + all_bbox[6][0]) / 2, (all_bbox[5][1] + all_bbox[6][1]) / 2,
        #                  (all_bbox[5][2] + all_bbox[6][2]) / 2, (all_bbox[5][3] + all_bbox[6][3]) / 2],
        #                 all_bbox[6],
        #                 [(all_bbox[6][0] + all_bbox[7][0]) / 2, (all_bbox[6][1] + all_bbox[7][1]) / 2,
        #                  (all_bbox[6][2] + all_bbox[7][2]) / 2, (all_bbox[6][3] + all_bbox[7][3]) / 2],
        #                 all_bbox[7],
        #                 [(all_bbox[7][0] + all_bbox[8][0]) / 2, (all_bbox[7][1] + all_bbox[8][1]) / 2,
        #                  (all_bbox[7][2] + all_bbox[8][2]) / 2, (all_bbox[7][3] + all_bbox[8][3]) / 2],
        #                 all_bbox[8],
        #                 [(all_bbox[8][0] + all_bbox[9][0]) / 2, (all_bbox[8][1] + all_bbox[9][1]) / 2,
        #                  (all_bbox[8][2] + all_bbox[9][2]) / 2, (all_bbox[8][3] + all_bbox[9][3]) / 2]]
        #
        #     # print('all_bbox:', all_bbox)
        #     # print('fram_order:', fram_order)
        #     # print('obs_traj:', obs_traj)
        #
        #     bbox = pred_bbox[:, best_idx] / scale_z
        #     lr = penalty[best_idx] * score[best_idx] * cfg.TRACK.LR
        #
        #     cx = bbox[0] + self.center_pos[0]
        #     cy = bbox[1] + self.center_pos[1]
        #
        #     # smooth bbox
        #     width = self.size[0] * (1 - lr) + bbox[2] * lr
        #     height = self.size[1] * (1 - lr) + bbox[3] * lr
        #
        #     # clip boundary
        #     cx, cy, width, height = self._bbox_clip(cx, cy, width, height, img.shape[:2])
        #
        #     # udpate state
        #     self.center_pos = np.array([cx, cy])
        #     self.size = np.array([width, height])
        #     bbox = [cx - width / 2, cy - height / 2, width, height]
        #     best_score = score[best_idx]
        #
        #     print('judge', 'bbox{}:'.format(fram_order), bbox)


        # elif threshold < 2 and block == 0 and all_len == 6:
        # elif score[best_idx] < 0.999999 and block == 0 and all_len >= 6 and all_len < 18:
        # # elif min_peak_distance < 5.5 and block == 0 and all_len == 6:
        #     print('*********************** 6 ** block ***********************')
        #     block_fram = 0
        #     block = 1
        #
        #     obs_traj = [all_bbox[all_len - 6], all_bbox[all_len - 6], all_bbox[all_len - 6],
        #                 all_bbox[all_len - 5], all_bbox[all_len - 5], all_bbox[all_len - 5],
        #                 all_bbox[all_len - 4], all_bbox[all_len - 4], all_bbox[all_len - 4],
        #                 all_bbox[all_len - 3], all_bbox[all_len - 3], all_bbox[all_len - 3],
        #                 all_bbox[all_len - 2], all_bbox[all_len - 2], all_bbox[all_len - 2],
        #                 all_bbox[all_len - 1], all_bbox[all_len - 1], all_bbox[all_len - 1] ]
        #
        #     bbox = pred_bbox[:, best_idx] / scale_z
        #     lr = penalty[best_idx] * score[best_idx] * cfg.TRACK.LR
        #
        #     cx = bbox[0] + self.center_pos[0]
        #     cy = bbox[1] + self.center_pos[1]
        #
        #     # smooth bbox
        #     width = self.size[0] * (1 - lr) + bbox[2] * lr
        #     height = self.size[1] * (1 - lr) + bbox[3] * lr
        #
        #     # clip boundary
        #     cx, cy, width, height = self._bbox_clip(cx, cy, width, height, img.shape[:2])
        #
        #     # udpate state
        #     self.center_pos = np.array([cx, cy])
        #     self.size = np.array([width, height])
        #     bbox = [cx - width / 2, cy - height / 2, width, height]
        #     best_score = score[best_idx]
        #
        #     print('judge', 'bbox{}:'.format(fram_order), bbox)

        # elif score[best_idx] < 0.999999 and block == 0 and all_len >= 3 and all_len < 6:
        # # elif min_peak_distance < 5.5 and block == 0 and all_len == 6:
        #     print('*********************** 6 ** block ***********************')
        #     block_fram = 0
        #     block = 1
        #
        #     obs_traj = [all_bbox[all_len - 3], all_bbox[all_len - 3], all_bbox[all_len - 3],
        #                 all_bbox[all_len - 3], all_bbox[all_len - 3], all_bbox[all_len - 3],
        #                 all_bbox[all_len - 2], all_bbox[all_len - 2], all_bbox[all_len - 2],
        #                 all_bbox[all_len - 2], all_bbox[all_len - 2], all_bbox[all_len - 2],
        #                 all_bbox[all_len - 1], all_bbox[all_len - 1], all_bbox[all_len - 1],
        #                 all_bbox[all_len - 1], all_bbox[all_len - 1], all_bbox[all_len - 1] ]
        #
        #     bbox = pred_bbox[:, best_idx] / scale_z
        #     lr = penalty[best_idx] * score[best_idx] * cfg.TRACK.LR
        #
        #     cx = bbox[0] + self.center_pos[0]
        #     cy = bbox[1] + self.center_pos[1]
        #
        #     # smooth bbox
        #     width = self.size[0] * (1 - lr) + bbox[2] * lr
        #     height = self.size[1] * (1 - lr) + bbox[3] * lr
        #
        #     # clip boundary
        #     cx, cy, width, height = self._bbox_clip(cx, cy, width, height, img.shape[:2])
        #
        #     # udpate state
        #     self.center_pos = np.array([cx, cy])
        #     self.size = np.array([width, height])
        #     bbox = [cx - width / 2, cy - height / 2, width, height]
        #     best_score = score[best_idx]
        #
        #     print('judge', 'bbox{}:'.format(fram_order), bbox)


        # test
        elif block == 1:  # prediction
            # print('obs_traj:', obs_traj)
            bbox = general_pred(obs_traj, block_fram)
            best_score = score[best_idx]
            block_fram = block_fram + 1
            if block_fram == 12:
                block = 0

            print('predict', 'bbox{}:'.format(fram_order), bbox)

        else:
            bbox = pred_bbox[:, best_idx] / scale_z
            lr = penalty[best_idx] * score[best_idx] * cfg.TRACK.LR

            cx = bbox[0] + self.center_pos[0]
            cy = bbox[1] + self.center_pos[1]

            # smooth bbox
            width = self.size[0] * (1 - lr) + bbox[2] * lr
            height = self.size[1] * (1 - lr) + bbox[3] * lr

            # clip boundary
            cx, cy, width, height = self._bbox_clip(cx, cy, width,
                                                    height, img.shape[:2])

            # udpate state
            self.center_pos = np.array([cx, cy])
            self.size = np.array([width, height])

            bbox = [cx - width / 2,
                    cy - height / 2,
                    width,
                    height]
            best_score = score[best_idx]

            # print('other', 'bbox{}:'.format(fram_order), bbox)

        # print('threshold{}:'.format(fram_order), threshold)
        # print('best', 'score{}:'.format(fram_order), best_score)
        # print('peak', 'distance{}:'.format(fram_order), min_peak_distance)

        return {
                'bbox': bbox,
                'best_score': best_score,
                'distance': peak_distance
               }
